import os
import csv
import subprocess
import copy


def read_csv(data_path):
    with open(data_path, "r") as f:
        data = csv.reader(f)
        list_data = list(data)

    return list_data


class OpProfiling(object):
    def __init__(self):
        self.bound_list = ["vec_bound", "mac_bound", "scalar_bound", "mte1_bound", "mte2_bound", "mte3_bound"]
        self.info_name = ["OP Type", "Task Duration(us)", "aicore_time(us)", "Block Dim", "vec_ratio", "mac_ratio",
                          "scalar_ratio", "mte1_ratio", "mte2_ratio", "mte3_ratio"]

    def judge_bound(self, ratios):
        bound_type = "no_bound"
        max_idx = ratios.index(max(ratios))
        if ratios[max_idx] > 0.8:
            bound_type = self.bound_list[max_idx]

        return bound_type

    def get_profiling_info(self, csv_data):
        info_idx = [csv_data[0].index(i) for i in self.info_name]

        res = []
        for d in csv_data[1:]:
            vec_ratio = round(float(d[info_idx[4]]), 2)
            mac_ratio = round(float(d[info_idx[5]]), 2)
            scalar_ratio = round(float(d[info_idx[6]]), 2)
            mte1_ratio = round(float(d[info_idx[7]]), 2)
            mte2_ratio = round(float(d[info_idx[8]]), 2)
            mte3_ratio = round(float(d[info_idx[9]]), 2)
            bound_type = self.judge_bound([vec_ratio, mac_ratio, scalar_ratio, mte1_ratio, mte2_ratio, mte3_ratio])

            res.append({"op": str(d[info_idx[0]]), "task_time": round(float(d[info_idx[1]]), 2),
                        "aicore_time": round(float(d[info_idx[2]]), 2),
                        "block_dim": int(d[info_idx[3]]), "vec_ratio": vec_ratio, "mac_ratio": mac_ratio,
                        "scalar_ratio": scalar_ratio, "mte1_ratio": mte1_ratio, "mte2_ratio": mte2_ratio,
                        "mte3_ratio": mte3_ratio, "bound_type": bound_type})

        return res

    def get_length(self, header_names, data):
        size = []
        for i in range(len(header_names)):
            col_data = [len(header_names[i])]
            for d in data:
                col_data.append(len(str(d[i])))
            size.append(max(col_data) + 4)

        return size

    def format_str(self, info, length):
        info_strs = []
        for i in range(len(info)):
            line_l = length[i]
            prof_data = info[i]
            info_str = "  " + str(prof_data) + " " * (line_l - len(str(prof_data)) - 2)
            info_strs.append(info_str)
        res = "|" + "|".join(info_strs) + "|"
        return res

    def format_print(self, prof_res):
        data = []
        for r in prof_res:
            data.append([r["op"], r["task_time"], r["aicore_time"], r["block_dim"], r["vec_ratio"], r["mac_ratio"],
                         r["scalar_ratio"], r["mte1_ratio"], r["mte2_ratio"], r["mte3_ratio"], r["bound_type"]])

        header_name = copy.deepcopy(self.info_name)
        header_name.append("bound_type")
        length = self.get_length(header_name, data)
        format_strs = [self.format_str(header_name, length)]
        for info in data:
            format_strs.append(self.format_str(info, length))

        base_line = ["-" * i for i in length]
        start_end_line = "+" + "+".join(base_line) + "+"

        print(start_end_line)
        for i in format_strs:
            print(i)
            print(start_end_line)

    def op_profiling(self, data_path, device_id=0):
        ms_path = os.environ["MSPROFPATH"]
        ms_py = "msprof.py"
        names = os.listdir(data_path)
        job_path = None
        res = []

        for name in names:
            if "PROF" in name or "JOB" in name:
                job_path = os.path.join(data_path, name)
                new_job_path = os.path.join(job_path, "device_%s" % device_id)
                if os.path.exists(new_job_path):
                    job_path = new_job_path
                break

        if os.path.exists(job_path):
            msprof_cmd = "python3 %s export summary -dir=%s" % (os.path.join(ms_path, ms_py), job_path)
            subprocess.run(msprof_cmd, shell=True)
            summary_path = os.path.join(job_path, "summary")
            if os.path.exists(summary_path):
                names = os.listdir(summary_path)
                result_name = None
                for name in names:
                    if "op_summary" in name:
                        result_name = name
                        break

                if not result_name:
                    return []
                csv_data = read_csv(os.path.join(summary_path, result_name))
                res = self.get_profiling_info(csv_data)

            return res


op_profiling = OpProfiling()
